package com.weaponlin.stream.collector;

import com.weaponlin.entity.Person;

import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collector;
import java.util.stream.Stream;

import static java.util.stream.Collector.Characteristics.CONCURRENT;
import static java.util.stream.Collector.Characteristics.IDENTITY_FINISH;
import static java.util.stream.Collectors.toMap;

public class MyGrouping implements Collector<Person, Map<Integer, ArrayList<Person>>, Map<Integer, ArrayList<Person>>> {
    @Override
    public Supplier<Map<Integer, ArrayList<Person>>> supplier() {
        return HashMap::new;
    }

    @Override
    public BiConsumer<Map<Integer, ArrayList<Person>>, Person> accumulator() {
        return (map, p) -> {
            ArrayList<Person> list;
            if ((list = map.get(p.getAge())) != null) {
                list.add(p);
            } else {
                list = new ArrayList<>();
                list.add(p);
                map.put(p.getAge(), list);
            }
        };
    }

    @Override
    public BinaryOperator<Map<Integer, ArrayList<Person>>> combiner() {
        return (m1, m2) -> Stream.of(m1, m2)
                .map(Map::entrySet)
                .flatMap(Collection::stream)
                .collect(toMap(Map.Entry::getKey, Map.Entry::getValue, (e1, e2) -> {
                    e1.addAll(e2);
                    return e1;
                }));
    }

    @Override
    public Function<Map<Integer, ArrayList<Person>>, Map<Integer, ArrayList<Person>>> finisher() {
        return Function.identity();
    }

    @Override
    public Set<Characteristics> characteristics() {
        return Collections.unmodifiableSet(EnumSet.of(IDENTITY_FINISH, CONCURRENT));
    }
}
